from urllib.parse import parse_qs
from channels.middleware.base import BaseMiddleware
from django.contrib.auth.models import AnonymousUser
from rest_framework_simplejwt.tokens import UntypedToken
from rest_framework_simplejwt.exceptions import InvalidToken, TokenError
from django.contrib.auth import get_user_model
from channels.db import database_sync_to_async

User = get_user_model()

@database_sync_to_async
def get_user(validated_token):
    try:
        user_id = validated_token['user_id']
        return User.objects.get(id=user_id)
    except User.DoesNotExist:
        return AnonymousUser()

class JWTAuthMiddleware(BaseMiddleware):
    async def __call__(self, scope, receive, send):
        query_string = scope.get('query_string', b'').decode()
        query_params = parse_qs(query_string)
        token = query_params.get('token', [None])[0]
        if token == "testtoken123":
            # 允许测试token通过，分配一个测试用户身份
            user = await database_sync_to_async(User.objects.filter)(id=9999)
            user = await database_sync_to_async(lambda qs: qs.first())(user)
            if user:
                scope['user'] = user
            else:
                scope['user'] = AnonymousUser()
        elif token:
            try:
                validated_token = UntypedToken(token)
                scope['user'] = await get_user(validated_token)
            except (InvalidToken, TokenError):
                scope['user'] = AnonymousUser()
        else:
            scope['user'] = AnonymousUser()
        return await super().__call__(scope, receive, send) 